import numpy as np
import cv2
import tritonclient.grpc as grpcclient
import time

if __name__ == '__main__':
    # 初始化 gRPC 客户端
    triton_client = grpcclient.InferenceServerClient(url='172.17.0.8:8001')  # gRPC 默认端口是 8001
    score_threshold = 0.3
    input_path = "/workspace/workspace/wumh/Motorcycle/dataset/images/test/img.png"
    # 读取图像
    image = cv2.imread(input_path)
    img = image.transpose((1, 0, 2))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    if img is None:
        raise FileNotFoundError(f"Image at path {input_path} not found")
    # 设置输入
    inputs = [
        grpcclient.InferInput('image', [*img.shape], "UINT8"),
        grpcclient.InferInput('score', [1], "FP16")
    ]
    inputs[0].set_data_from_numpy(img)
    inputs[1].set_data_from_numpy(np.array([score_threshold], dtype=np.float16))

    # 设置输出
    outputs = [
        grpcclient.InferRequestedOutput('scores'),
        grpcclient.InferRequestedOutput('bboxes'),
    ]

    t1 = time.time()
    infer_result = triton_client.infer('base', inputs=inputs, outputs=outputs)
    t2 = time.time()

    # 获取推理结果
    bboxes = infer_result.as_numpy('bboxes')
    scores = infer_result.as_numpy('scores')

    for i in range(len(bboxes)):
        print(f"score: [{round(scores[i], 4)}]    bbox: {bboxes[i]}")

    print('inference time is: {}ms'.format(1000 * (t2 - t1)))
